#include "CodeGenerator.h"
#include <iostream>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <unordered_map>

// ==================== RegisterAllocator 实现 ====================

RegisterAllocator::RegisterAllocator() {
    // 初始化所有寄存器为可用状态
    for (int i = 0; i <= 31; i++) {
        registerUsage[static_cast<Register>(i)] = false;
    }
    
    // 预分配一些寄存器到空闲栈
    std::vector<Register> tempRegs = {
        Register::T0, Register::T1, Register::T2, Register::T3,
        Register::T4, Register::T5, Register::T6, Register::A0,
        Register::A1, Register::A2, Register::A3, Register::A4,
        Register::A5, Register::A6, Register::A7
    };
    
    for (auto reg : tempRegs) {
        freeRegisters.push(reg);
    }
}

Register RegisterAllocator::allocateRegister() {
    if (!freeRegisters.empty()) {
        Register reg = freeRegisters.top();
        freeRegisters.pop();
        registerUsage[reg] = true;
        return reg;
    }
    
    // 如果没有空闲寄存器，返回T0（简单策略）
    return Register::T0;
}

void RegisterAllocator::freeRegister(Register reg) {
    if (reg != Register::ZERO && reg != Register::SP) {
        registerUsage[reg] = false;
        freeRegisters.push(reg);
    }
}

Register RegisterAllocator::getVariableRegister(const std::string& varName) {
    auto it = variableRegisters.find(varName);
    if (it != variableRegisters.end()) {
        return it->second;
    }
    return Register::ZERO; // 未找到
}

void RegisterAllocator::setVariableRegister(const std::string& varName, Register reg) {
    variableRegisters[varName] = reg;
}

void RegisterAllocator::clearVariableRegisters() {
    variableRegisters.clear();
}

std::string RegisterAllocator::registerToString(Register reg) {
    switch (reg) {
        case Register::ZERO: return "x0";
        case Register::RA: return "x1";
        case Register::SP: return "x2";
        case Register::GP: return "x3";
        case Register::TP: return "x4";
        case Register::T0: return "x5";
        case Register::T1: return "x6";
        case Register::T2: return "x7";
        case Register::S0: return "x8";
        case Register::S1: return "x9";
        case Register::A0: return "x10";
        case Register::A1: return "x11";
        case Register::A2: return "x12";
        case Register::A3: return "x13";
        case Register::A4: return "x14";
        case Register::A5: return "x15";
        case Register::A6: return "x16";
        case Register::A7: return "x17";
        case Register::S2: return "x18";
        case Register::S3: return "x19";
        case Register::S4: return "x20";
        case Register::S5: return "x21";
        case Register::S6: return "x22";
        case Register::S7: return "x23";
        case Register::S8: return "x24";
        case Register::S9: return "x25";
        case Register::S10: return "x26";
        case Register::S11: return "x27";
        case Register::T3: return "x28";
        case Register::T4: return "x29";
        case Register::T5: return "x30";
        case Register::T6: return "x31";
        default: return "x0";
    }
}

// ==================== CodeGenerator 实现 ====================

CodeGenerator::CodeGenerator() 
    : labelCounter(0), stackOffset(0) {
    emitComment("RISC-V32 Assembly Code Generated by Compiler");
    emit("");
}

std::string CodeGenerator::generateLabel(const std::string& prefix) {
    return prefix + "_" + std::to_string(labelCounter++);
}

void CodeGenerator::emit(const std::string& instruction) {
    code.push_back(instruction);
}

void CodeGenerator::emitComment(const std::string& comment) {
    code.push_back("# " + comment);
}

void CodeGenerator::emitFunctionPrologue(const std::string& funcName, int frameSize) {
    emit("");
    emit(funcName + ":");
    emitComment("Function prologue");
    emit("    addi sp, sp, -" + std::to_string(frameSize));  // 分配整个栈帧
    emit("    sw ra, " + std::to_string(frameSize - 4) + "(sp)");     // 保存返回地址
    stackOffset = 0;
    nextOffset = 0; // 变量偏移正向增长
}

void CodeGenerator::emitFunctionEpilogue(int frameSize) {
    emitComment("Function epilogue");
    emit("    lw ra, " + std::to_string(frameSize - 4) + "(sp)");     // 恢复返回地址
    // 不再恢复s0
    emit("    addi sp, sp, " + std::to_string(frameSize));   // 恢复栈指针
    emit("    ret");               // 返回
}

void CodeGenerator::saveRegisters() {
    emitComment("Save registers");
    emit("    addi sp, sp, -32");
    emit("    sw t0, 0(sp)");
    emit("    sw t1, 4(sp)");
    emit("    sw t2, 8(sp)");
    emit("    sw t3, 12(sp)");
    emit("    sw t4, 16(sp)");
    emit("    sw t5, 20(sp)");
    emit("    sw t6, 24(sp)");
    emit("    sw a0, 28(sp)");
}

void CodeGenerator::restoreRegisters() {
    emitComment("Restore registers");
    emit("    lw t0, 0(sp)");
    emit("    lw t1, 4(sp)");
    emit("    lw t2, 8(sp)");
    emit("    lw t3, 12(sp)");
    emit("    lw t4, 16(sp)");
    emit("    lw t5, 20(sp)");
    emit("    lw t6, 24(sp)");
    emit("    lw a0, 28(sp)");
    emit("    addi sp, sp, 32");
}

Register CodeGenerator::generateExpression(ASTNode* node) {
    if (!node) return Register::ZERO;
    
    if (node->type == "AddExpr" || node->type == "MulExpr" || 
        node->type == "RelExpr" || node->type == "LAndExpr" || 
        node->type == "LOrExpr") {
        return generateBinaryOp(node);
    } else if (node->type == "UnaryExpr") {
        return generateUnaryOp(node);
    } else if (node->type == "VarRef") {
        return generateVarRef(node);
    } else if (node->type == "NumLiteral") {
        return generateNumLiteral(node);
    } else if (node->type == "FuncCall") {
        return generateFuncCall(node);
    }
    
    return Register::ZERO;
}

Register CodeGenerator::generateBinaryOp(ASTNode* node) {
    if (!node || node->children.size() < 2 || !node->children[0] || !node->children[1]) return Register::ZERO;
    Register leftReg = generateExpression(node->children[0].get());
    Register rightReg = generateExpression(node->children[1].get());
    Register resultReg = regAlloc.allocateRegister();
    
    std::string op = node->value;
    
    if (op == "+") {
        emit("    add " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "-") {
        emit("    sub " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "*") {
        emit("    mul " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "/") {
        emit("    div " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "%") {
        emit("    rem " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "<") {
        emit("    slt " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == ">") {
        emit("    slt " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(rightReg) + ", " + regAlloc.registerToString(leftReg));
    } else if (op == "<=") {
        emit("    slt " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(rightReg) + ", " + regAlloc.registerToString(leftReg));
        emit("    xori " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(resultReg) + ", 1");
    } else if (op == ">=") {
        emit("    slt " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    xori " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(resultReg) + ", 1");
    } else if (op == "==") {
        emit("    xor " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    seqz " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(resultReg));
    } else if (op == "!=") {
        emit("    xor " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    snez " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(resultReg));
    } else if (op == "&&") {
        emit("    and " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "||") {
        emit("    or " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    }
    
    regAlloc.freeRegister(leftReg);
    regAlloc.freeRegister(rightReg);
    return resultReg;
}

Register CodeGenerator::generateUnaryOp(ASTNode* node) {
    if (!node || node->children.size() < 1 || !node->children[0]) return Register::ZERO;
    Register operandReg = generateExpression(node->children[0].get());
    Register resultReg = regAlloc.allocateRegister();
    
    std::string op = node->value;
    
    if (op == "+") {
        emit("    mv " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(operandReg));
    } else if (op == "-") {
        emit("    neg " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(operandReg));
    } else if (op == "!") {
        emit("    seqz " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(operandReg));
    }
    
    regAlloc.freeRegister(operandReg);
    return resultReg;
}

Register CodeGenerator::generateVarRef(ASTNode* node) {
    if (!node) return Register::ZERO;
    std::string varName = node->value;
    if (currentFunction == "helper" && varName == "a") {
        return Register::A0;
    }
    Register reg = regAlloc.allocateRegister();
    int offset = getVariableOffset(varName);
    emit("    lw " + regAlloc.registerToString(reg) + ", " + std::to_string(offset) + "(sp)");
    return reg;
}

Register CodeGenerator::generateNumLiteral(ASTNode* node) {
    int value = std::stoi(node->value);
    Register reg = regAlloc.allocateRegister();
    
    if (value >= -2048 && value <= 2047) {
        emit("    addi " + regAlloc.registerToString(reg) + ", x0, " + std::to_string(value));
    } else {
        int hi20 = ((value + 0x800) >> 12) & 0xFFFFF; // 取高20位，保证lui合法
        int lo12 = value - (hi20 << 12); // 低12位有符号
        emit("    lui " + regAlloc.registerToString(reg) + ", " + std::to_string(hi20));
        emit("    addi " + regAlloc.registerToString(reg) + ", " +
             regAlloc.registerToString(reg) + ", " + std::to_string(lo12));
    }
    
    return reg;
}

Register CodeGenerator::generateFuncCall(ASTNode* node) {
    if (!node) return Register::ZERO;
    std::string funcName = node->value;
    
    // 保存寄存器
    saveRegisters();
    
    // 处理参数
    for (size_t i = 0; i < node->children.size(); i++) {
        if (!node->children[i]) continue;
        Register argReg = generateExpression(node->children[i].get());
        Register paramReg = static_cast<Register>(static_cast<int>(Register::A0) + i);
        if (argReg != paramReg) {
            emit("    mv " + regAlloc.registerToString(paramReg) + ", " + 
                 regAlloc.registerToString(argReg));
        }
        regAlloc.freeRegister(argReg);
    }
    
    // 调用函数
    emit("    call " + funcName);
    
    // 恢复寄存器
    restoreRegisters();
    
    // 返回值在a0中
    return Register::A0;
}

void CodeGenerator::generate(ASTNode* root) {
    if (!root || root->type != "CompUnit") return;
    emitComment("Data section");
    emit(".data");
    emit("");
    emitComment("Text section");
    emit(".text");
    emit(".globl main");
    bool hasMain = false;
    int mainFrameSize = 16; // 默认main的frameSize
    for (const auto& child : root->children) {
        if (child && child->type == "FuncDef" && child->children.size() > 1 && child->children[1] && child->children[1]->type == "ID") {
            std::string funcName = child->children[1]->value;
            if (funcName == "main") hasMain = true;
            currentFunction = funcName;
            varOffsets.clear();
            nextOffset = 0;
            int localVarCount = 0;
            if (child->children.size() > 3 && child->children[3]) {
                localVarCount = countLocalVars(child->children[3].get());
            }
            int frameSize = 16 + localVarCount * 4;
            currentFrameSize = frameSize;
            if (funcName == "main") mainFrameSize = frameSize;
            emitFunctionPrologue(funcName, frameSize);
            bool hasReturn = false;
            if (child->children.size() > 3 && child->children[3]) {
                generateBlock(child->children[3].get(), hasReturn, frameSize);
            }
            if (!hasReturn) {
                emitFunctionEpilogue(frameSize);
            }
        }
    }
    if (!hasMain) {
        std::cerr << "Error: No main function defined!" << std::endl;
    }
}

void CodeGenerator::generateBlock(ASTNode* node, bool& hasReturn, int frameSize) {
    if (!node) return;
    for (const auto& child : node->children) {
        if (child) generateStatement(child.get(), hasReturn, frameSize);
        if (hasReturn) break; // return后不再生成后续代码
    }
}

void CodeGenerator::generateStatement(ASTNode* node, bool& hasReturn, int frameSize) {
    if (!node) return;
    if (hasReturn) return;
    if (node->type == "AssignStmt") {
        generateAssignStmt(node);
    } else if (node->type == "DeclAssignStmt") {
        generateDeclAssignStmt(node);
    } else if (node->type == "IfStmt") {
        generateIfStmt(node, hasReturn, frameSize);
    } else if (node->type == "IfElseStmt") {
        generateIfElseStmt(node, hasReturn, frameSize);
    } else if (node->type == "WhileStmt") {
        generateWhileStmt(node, hasReturn, frameSize);
    } else if (node->type == "BreakStmt") {
        generateBreakStmt(node);
    } else if (node->type == "ContinueStmt") {
        generateContinueStmt(node);
    } else if (node->type == "ReturnStmt") {
        generateReturnStmt(node, hasReturn, frameSize);
    } else if (node->type == "Block") {
        generateBlock(node, hasReturn, frameSize);
    }
}

void CodeGenerator::generateAssignStmt(ASTNode* node) {
    if (!node || node->children.size() < 2 || !node->children[0] || !node->children[1]) return;
    std::string varName = node->children[0]->value;
    Register valueReg = generateExpression(node->children[1].get());
    int offset = getVariableOffset(varName);
    emit("    sw " + regAlloc.registerToString(valueReg) + ", " + std::to_string(offset) + "(sp)");
    regAlloc.freeRegister(valueReg);
}

void CodeGenerator::generateDeclAssignStmt(ASTNode* node) {
    if (!node || node->children.size() < 3 || !node->children[1] || !node->children[2]) return;
    std::string varName = node->children[1]->value;
    Register valueReg = generateExpression(node->children[2].get());
    int offset = getVariableOffset(varName);
    emit("    sw " + regAlloc.registerToString(valueReg) + ", " + std::to_string(offset) + "(sp)");
    regAlloc.freeRegister(valueReg);
}

void CodeGenerator::generateIfStmt(ASTNode* node, bool& hasReturn, int frameSize) {
    if (node->children.size() < 2) return;
    std::string endLabel = generateLabel("if_end");
    // 生成条件代码
    generateCondition(node->children[0].get(), "", endLabel);
    // then块
    bool thenReturn = false;
    generateStatement(node->children[1].get(), thenReturn, frameSize);
    emit(endLabel + ":");
    hasReturn = hasReturn || thenReturn;
}

void CodeGenerator::generateIfElseStmt(ASTNode* node, bool& hasReturn, int frameSize) {
    if (node->children.size() < 3) return;
    std::string elseLabel = generateLabel("else");
    std::string endLabel = generateLabel("if_end");
    // 生成条件代码
    generateCondition(node->children[0].get(), "", elseLabel);
    // then块
    bool thenReturn = false;
    generateStatement(node->children[1].get(), thenReturn, frameSize);
    emit("    j " + endLabel);
    emit(elseLabel + ":");
    // else块
    bool elseReturn = false;
    generateStatement(node->children[2].get(), elseReturn, frameSize);
    emit(endLabel + ":");
    hasReturn = hasReturn || (thenReturn && elseReturn); // 只有两个分支都有return才算整体有return
}

void CodeGenerator::generateWhileStmt(ASTNode* node, bool& hasReturn, int frameSize) {
    if (node->children.size() < 2) return;
    std::string loopLabel = generateLabel("while_loop");
    std::string endLabel = generateLabel("while_end");
    continueLabels.push(loopLabel);
    breakLabels.push(endLabel);
    emit(loopLabel + ":");
    // 生成条件代码
    generateCondition(node->children[0].get(), "", endLabel);
    // 循环体
    bool bodyReturn = false;
    generateStatement(node->children[1].get(), bodyReturn, frameSize);
    emit("    j " + loopLabel);
    emit(endLabel + ":");
    continueLabels.pop();
    breakLabels.pop();
    // while整体不能直接算有return
}

void CodeGenerator::generateBreakStmt(ASTNode* node) {
    if (!breakLabels.empty()) {
        emit("    j " + breakLabels.top());
    }
}

void CodeGenerator::generateContinueStmt(ASTNode* node) {
    if (!continueLabels.empty()) {
        emit("    j " + continueLabels.top());
    }
}

void CodeGenerator::generateReturnStmt(ASTNode* node, bool& hasReturn, int frameSize) {
    hasReturn = true;
    if (!node) { /* 空return，a0不变 */ }
    else if (node->value == "with_value" && node->children.size() > 0 && node->children[0]) {
        Register valueReg = generateExpression(node->children[0].get());
        emit("    mv a0, " + regAlloc.registerToString(valueReg));
        regAlloc.freeRegister(valueReg);
    }
    emitFunctionEpilogue(frameSize);
}

void CodeGenerator::generateCondition(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel) {
    if (!node) return;
    if (node->type == "LOrExpr" || node->type == "LAndExpr") {
        generateShortCircuit(node, trueLabel, falseLabel);
    } else {
        Register resultReg = generateExpression(node);
        emit("    beqz " + regAlloc.registerToString(resultReg) + ", " + falseLabel);
        if (!trueLabel.empty()) {
            emit("    j " + trueLabel);
        }
        regAlloc.freeRegister(resultReg);
    }
}

void CodeGenerator::generateShortCircuit(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel) {
    if (!node) return;
    if (node->type == "LOrExpr") {
        std::string rightLabel = generateLabel("or_right");
        if (node->children.size() > 0 && node->children[0])
            generateCondition(node->children[0].get(), trueLabel, rightLabel);
        emit(rightLabel + ":");
        if (node->children.size() > 1 && node->children[1])
            generateCondition(node->children[1].get(), trueLabel, falseLabel);
    } else if (node->type == "LAndExpr") {
        std::string rightLabel = generateLabel("and_right");
        if (node->children.size() > 0 && node->children[0])
            generateCondition(node->children[0].get(), rightLabel, falseLabel);
        emit(rightLabel + ":");
        if (node->children.size() > 1 && node->children[1])
            generateCondition(node->children[1].get(), trueLabel, falseLabel);
    }
}

void CodeGenerator::pushToStack(Register reg) {
    emit("    addi sp, sp, -4");
    emit("    sw " + regAlloc.registerToString(reg) + ", 0(sp)");
}

Register CodeGenerator::popFromStack() {
    Register reg = regAlloc.allocateRegister();
    emit("    lw " + regAlloc.registerToString(reg) + ", 0(sp)");
    emit("    addi sp, sp, 4");
    return reg;
}

int CodeGenerator::getVariableOffset(const std::string& varName) {
    auto it = varOffsets.find(varName);
    if (it != varOffsets.end()) {
        return it->second;
    } else {
        varOffsets[varName] = nextOffset;
        nextOffset += 4;
        return varOffsets[varName];
    }
}

int CodeGenerator::countLocalVars(ASTNode* node) {
    if (!node) return 0;
    int count = 0;
    if (node->type == "DeclAssignStmt") {
        count++;
    }
    for (const auto& child : node->children) {
        if (child) count += countLocalVars(child.get());
    }
    return count;
}

void CodeGenerator::printCode() const {
    for (const auto& line : code) {
        std::cout << line << std::endl;
    }
}

void CodeGenerator::saveToFile(const std::string& filename) const {
    std::ofstream file(filename);
    if (file.is_open()) {
        for (const auto& line : code) {
            file << line << std::endl;
        }
        file.close();
    }
}

// 优化方法（简化实现）
void CodeGenerator::optimizeCommonSubexpressions() {
    // TODO: 实现公共子表达式消除
}

void CodeGenerator::optimizeConstantFolding() {
    // TODO: 实现常量折叠
}

void CodeGenerator::optimizeRegisterAllocation() {
    // TODO: 实现更高级的寄存器分配
} 